MODULE md1_FE_AO_overlap

   USE md1_global_consts
   USE md1_global_types
   IMPLICIT NONE
   PRIVATE
   ! Public procedures
   PUBLIC :: md1_get_FE_AO_overlaps,        &
             md1_overlap_AO_FE

   ! Diagonostic variables only (not needed for calculation)
   REAL(REALK), SAVE, PUBLIC :: molsize_lo(3) = 1d10,          &
                                molsize_hi(3) = -1d10

CONTAINS

!-------------------------------------------------------------------------------

   SUBROUTINE md1_get_FE_AO_overlaps(basis,sh_pairs,grid,Q_matrix)

      USE md1_utils,      ONLY: md1_get_prim_batch
      USE md1_LIP_coeffs, ONLY: md1_init_LIP_coeffs, md1_free_LIP_coeffs

      IMPLICIT NONE
      TYPE(md1_basis),    INTENT(IN)  :: basis
      TYPE(md1_sh_pairs), INTENT(IN)  :: sh_pairs(:)
      TYPE(md1_FE_grid),  INTENT(IN)  :: grid
      REAL(REALK),        INTENT(OUT) :: Q_matrix(:,:)

      TYPE(md1_prim_batch) :: batch(basis%Cmax2)
      INTEGER(INTK) :: k, np

      CALL md1_init_LIP_coeffs(grid%order,grid%h)

      DO k = 1, SIZE(sh_pairs)
         CALL md1_get_prim_batch(basis,sh_pairs(k),batch,np)
         CALL md1_overlap_AO_FE(basis%LM,sh_pairs(k),batch(:np),grid,Q_matrix)
      END DO

      CALL md1_free_LIP_coeffs

   END SUBROUTINE md1_get_FE_AO_overlaps

!-------------------------------------------------------------------------------

   SUBROUTINE md1_overlap_AO_FE(LM,sh,batch,grid,Qmat)

      USE md1_hermite_finite_mpoles, ONLY: md1_build_1D_finite_mpoles

      IMPLICIT NONE
      TYPE(md1_basis_LM),   INTENT(IN)  :: LM(0:)
      TYPE(md1_sh_pairs),   INTENT(IN)  :: sh
      TYPE(md1_prim_batch), INTENT(IN)  :: batch(:)
      TYPE(md1_FE_grid),    INTENT(IN)  :: grid
      REAL(REALK),          INTENT(OUT) :: Qmat(:,:)

      REAL(REALK), POINTER :: Sptr(:,:,:,:)
      TYPE(md1_5d_array) :: Sjie(3)
      TYPE(md1_int_hilo) :: kFE(3)
      INTEGER(INTK) :: ndim, i,k, xyz, np
      REAL(REALK) :: abc(3)

      np = SIZE(batch)

      ! Get FEs overlapping with these Gaussian functions
      DO xyz = 1, 3
         CALL md1_get_1D_FE_range(xyz,batch,grid,kFE(xyz))
         ndim = kFE(xyz)%hi - kFE(xyz)%lo + 1
         ALLOCATE (Sjie(xyz)%a(np,0:sh%Jang,0:sh%Iang,0:grid%order,ndim))
      END DO

      DO xyz = 1, 3
         i = 0
         DO k = kFE(xyz)%lo, kFE(xyz)%hi
            i = i + 1
            ! Define integration range [a,b]
            abc(1) = grid%lo(xyz) + DBLE(k-1) * grid%L(xyz)
            abc(2) = grid%lo(xyz) + DBLE(k) * grid%L(xyz)
            ! Define multipole origin c
            abc(3) = abc(1)
            Sptr => Sjie(xyz)%a(:,:,:,:,i)

            ! Build 1D multipole integrals for this FE
            CALL md1_build_1D_finite_mpoles(xyz,batch,abc,Sptr)
            ! Transform e-index from monomials to FE polynomials
            CALL md1_transform_to_FE_overlaps(xyz,grid%order,Sptr)
         END DO
      END DO

      ! Contract 1D overlaps with 3D grid
      CALL md1_contract_with_Vpot(kFE,Sjie(1)%a,Sjie(2)%a,Sjie(3)%a,     &
                                  sh,LM,batch,(sh%I==sh%J),grid%f,Qmat)

      DO xyz = 1,3; DEALLOCATE (Sjie(xyz)%a); END DO

   END SUBROUTINE md1_overlap_AO_FE

!-------------------------------------------------------------------------------

   SUBROUTINE md1_get_1D_FE_range(xyz,batch,grid,kFE)

      IMPLICIT NONE
      INTEGER(INTK),        INTENT(IN)  :: xyz
      TYPE(md1_prim_batch), INTENT(IN)  :: batch(:)
      TYPE(md1_FE_grid),    INTENT(IN)  :: grid
      TYPE(md1_int_hilo),   INTENT(OUT) :: kFE

      REAL(REALK), PARAMETER :: thresh = 1d-10
      INTEGER(INTK) :: i
      REAL(REALK) :: r, lo, hi

      lo =  1d10
      hi = -1d10
      DO i = 1, SIZE(batch)
         r  = SQRT(-LOG(thresh) / batch(i)%exp)
         lo = MIN(lo, batch(i)%P(xyz) - r)
         hi = MAX(hi, batch(i)%P(xyz) + r)
      END DO

      IF (lo < grid%lo(xyz)) THEN 
         WRITE(IO1,*) lo, MINVAL(batch(:)%exp); STOP 'fixme grid lo size'
      END IF
      IF (hi > grid%hi(xyz)) THEN
         WRITE(IO1,*) hi, MINVAL(batch(:)%exp); STOP 'fixme grid hi size'
      END IF

      kFE%lo = CEILING( (lo-grid%lo(xyz)) /grid%L(xyz) )
      kFE%hi = CEILING( (hi-grid%lo(xyz)) /grid%L(xyz) )

      ! Diagnostic only
      molsize_lo(xyz) = MIN(molsize_lo(xyz), lo)
      molsize_hi(xyz) = MAX(molsize_hi(xyz), hi)

   END SUBROUTINE md1_get_1D_FE_range

!-------------------------------------------------------------------------------

   SUBROUTINE md1_transform_to_FE_overlaps(xyz,order,Sjie)

      USE md1_LIP_coeffs, ONLY: md1_get_LIP_coeffs
      
      IMPLICIT NONE
      INTEGER(INTK), INTENT(IN)    :: xyz, order
      REAL(REALK),   INTENT(INOUT) :: Sjie(:,:,:,:)
      
      REAL(REALK), POINTER :: coeffs(:,:)
      INTEGER(INTK) :: i,j,p
      
      CALL md1_get_LIP_coeffs(xyz,order,coeffs) 
      
      DO i = 1, SIZE(Sjie,3)
         DO j = 1, SIZE(Sjie,2)
            DO p = 1, SIZE(Sjie,1)
               Sjie(p,j,i,:) = MATMUL( coeffs(:,:), Sjie(p,j,i,:) )
            END DO
         END DO
      END DO
   
   END SUBROUTINE md1_transform_to_FE_overlaps

!-------------------------------------------------------------------------------

   SUBROUTINE md1_contract_with_Vpot(kFE,SX,SY,SZ,sh,LM,batch,IeqJ,Vpot,Qmat)

      IMPLICIT NONE
      TYPE(md1_int_hilo),   INTENT(IN)  :: kFE(3)
      REAL(REALK),  TARGET, INTENT(IN)  :: SX(:,0:,0:,0:,kFE(1)%lo:)
      REAL(REALK),  TARGET, INTENT(IN)  :: SY(:,0:,0:,0:,kFE(2)%lo:)
      REAL(REALK),  TARGET, INTENT(IN)  :: SZ(:,0:,0:,0:,kFE(3)%lo:)
      TYPE(md1_sh_pairs),   INTENT(IN)  :: sh
      TYPE(md1_basis_LM),   INTENT(IN)  :: LM(0:)
      TYPE(md1_prim_batch), INTENT(IN)  :: batch(:)
      LOGICAL,              INTENT(IN)  :: IeqJ
      REAL(REALK),          INTENT(IN)  :: Vpot(:,:,:)
      REAL(REALK),          INTENT(OUT) :: Qmat(:,:)

      REAL(REALK), POINTER :: X(:), Y(:), Z(:)
      INTEGER(INTK) :: ituv(3), jtuv(3)
      INTEGER(INTK) :: P,Q, pM,qM,qMhi
      INTEGER(INTK) :: i,j,k, ii,jj,kk, u,v,w, order
      REAL(REALK)   :: C

      order = UBOUND(SX,4)

      P = sh%P%lo
      DO pM = 1, LM(sh%Iang)%car_nlm
         ituv(:) = LM(sh%Iang)%tuv(:,pM)
         qMhi = LM(sh%Jang)%car_nlm
         IF (IeqJ) qMhi = pM
         Q = sh%Q%lo
         DO qM = 1, qMhi
            jtuv(:) = LM(sh%Jang)%tuv(:,qM)
            Qmat(P,Q) = zero

            ! Pick out overlapping FEs
            DO ii = kFE(1)%lo, kFE(1)%hi
               DO jj = kFE(2)%lo, kFE(2)%hi
                  DO kk = kFE(3)%lo, kFE(3)%hi

                     ! Grid points at FE corner
                     i = (ii-1) * order + 1
                     j = (jj-1) * order + 1
                     k = (kk-1) * order + 1

                     DO u = 0, order
                        DO v = 0, order
                           DO w = 0, order

!FIXME for speed, re-order matrix indices before calling these loops?!
                              X => SX(:,jtuv(1),ituv(1),u,ii)
                              Y => SY(:,jtuv(2),ituv(2),v,jj)
                              Z => SZ(:,jtuv(3),ituv(3),w,kk)
                              C = SUM( X(:)*Y(:)*Z(:) * batch(:)%CAB ) 
                              Qmat(P,Q) = Qmat(P,Q) + C * Vpot(i+u,j+v,k+w)

                           END DO
                        END DO
                     END DO

                  END DO
               END DO
            END DO

            IF ( P /= Q ) Qmat(Q,P) = Qmat(P,Q)
            Q = Q+1
         END DO
         P = P+1
      END DO

   END SUBROUTINE md1_contract_with_Vpot

!-------------------------------------------------------------------------------

END MODULE md1_FE_AO_overlap
